{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YadOCCCyXT0y"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "S7vZFQkeq_Vk"
      },
      "source": [
        "Clone threestudio repo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubuVj4z0MhHh"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/threestudio-project/threestudio.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HrGD3RtXXsB9"
      },
      "outputs": [],
      "source": [
        "%cd threestudio"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "aXqmxXX0Jb6m"
      },
      "source": [
        "Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HKZQNbNmX20t"
      },
      "outputs": [],
      "source": [
        "!pip install ninja\n",
        "!pip install lightning==2.0.0 omegaconf==2.3.0 jaxtyping typeguard diffusers transformers accelerate opencv-python tensorboard matplotlib imageio imageio[ffmpeg] trimesh bitsandbytes sentencepiece safetensors huggingface_hub libigl xatlas networkx pysdf PyMCubes wandb torchmetrics controlnet_aux\n",
        "!pip install einops kornia taming-transformers-rom1504 git+https://github.com/openai/CLIP.git # zero123\n",
        "!pip install open3d plotly # mesh visualization"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "b8U3rne0JTgs"
      },
      "source": [
        "And build some dependencies manually. This may take a while."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BFHI67e0_65a"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://github.com/ashawkey/envlight.git\n",
        "!pip install git+https://github.com/KAIR-BAIR/nerfacc.git@v0.5.2\n",
        "!pip install git+https://github.com/NVlabs/nvdiffrast.git\n",
        "!pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "QxY_7clwrJ18"
      },
      "source": [
        "Login to HuggingFace"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T4eeJFCZAZUV"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import interpreter_login\n",
        "\n",
        "interpreter_login()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Eeqa6u5QrPTj"
      },
      "source": [
        "Now create your own 3D model from text prompts\n",
        "\n",
        "Here we use the DreamFusion model with DeepFloyd-IF guidance. You may try other models by using different running commands given [here](https://github.com/threestudio-project/threestudio#supported-models)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pZUD5VcS88yz"
      },
      "outputs": [],
      "source": [
        "prompt = \"a zoomed out DSLR photo of a baby bunny sitting on top of a stack of pancakes\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uy61bM84K7Qi"
      },
      "outputs": [],
      "source": [
        "!python launch.py --config configs/dreamfusion-if.yaml --train --gpu 0 system.prompt_processor.prompt=\"$prompt\" trainer.max_steps=10000 system.prompt_processor.spawn=false"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "UqrrK9FxrYgd"
      },
      "source": [
        "Display the rendered video"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J14FrBwMk5m9"
      },
      "outputs": [],
      "source": [
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "def display_video(video_path):\n",
        "  mp4 = open(video_path,'rb').read()\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "  return HTML(\"\"\"\n",
        "  <video width=1000 controls>\n",
        "    <source src=\"%s\" type=\"video/mp4\">\n",
        "  </video>\n",
        "  \"\"\" % data_url)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aaQCVmiCp87c"
      },
      "outputs": [],
      "source": [
        "# you will see the path to the saving directory at the end of the training logs\n",
        "# replace save_dir below with that path\n",
        "save_dir = 'path/to/save/dir'\n",
        "\n",
        "import os\n",
        "import glob\n",
        "video_path = glob.glob(os.path.join(save_dir, \"*-test.mp4\"))[0]\n",
        "display_video(video_path)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "UlUL4ZPn2xQV"
      },
      "source": [
        "Extract the object mesh.\n",
        "\n",
        "Here we use an empirical threshold value. You can also first try `system.geometry.isosurface_threshold=auto` and visualize it. Then you can manually adjust the threshold according to the automatically determined value shown in the logs. Increase it if there are too many floaters and decrease it if the geometry is incomplete. \n",
        "\n",
        "\n",
        "The extraction process takes around 2 mins on T4."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f-R26yEXqgyL"
      },
      "outputs": [],
      "source": [
        "!python launch.py --config $save_dir/../configs/parsed.yaml --export --gpu 0 resume=$save_dir/../ckpts/last.ckpt system.exporter_type=mesh-exporter system.exporter.context_type=cuda system.geometry.isosurface_threshold=15.0 "
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "iocD0gl870vq"
      },
      "source": [
        "Visualize the mesh. Or you can directly download the export assets and use them locally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "71-J9PRIyp99"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import open3d as o3d\n",
        "import plotly.graph_objects as go\n",
        "import glob\n",
        "\n",
        "mesh_path = glob.glob(os.path.join(save_dir, \"*-export/model.obj\"))[0]\n",
        "mesh = o3d.io.read_triangle_mesh(mesh_path)\n",
        "if not mesh.has_vertex_normals():\n",
        "  mesh.compute_vertex_normals()\n",
        "if not mesh.has_triangle_normals():\n",
        "  mesh.compute_triangle_normals()\n",
        "\n",
        "triangles = np.asarray(mesh.triangles)\n",
        "vertices = np.asarray(mesh.vertices)\n",
        "colors = None\n",
        "if mesh.has_triangle_normals():\n",
        "  colors = (0.5, 0.5, 0.5) + np.asarray(mesh.triangle_normals) * 0.5\n",
        "  colors = tuple(map(tuple, colors))\n",
        "else:\n",
        "  colors = (1.0, 0.0, 0.0)\n",
        "fig = go.Figure(\n",
        "  data=[\n",
        "    go.Mesh3d(\n",
        "      x=vertices[:,0],\n",
        "      y=vertices[:,1],\n",
        "      z=vertices[:,2],\n",
        "      i=triangles[:,0],\n",
        "      j=triangles[:,1],\n",
        "      k=triangles[:,2],\n",
        "      facecolor=colors,\n",
        "      opacity=0.50)\n",
        "  ],\n",
        "  layout=dict(\n",
        "    scene=dict(\n",
        "      xaxis=dict(visible=False),\n",
        "      yaxis=dict(visible=False),\n",
        "      zaxis=dict(visible=False)\n",
        "    )\n",
        "  )\n",
        ")\n",
        "fig.show()\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
